package cc.springai.langchain4j.controller;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingLanguageModel;
import dev.langchain4j.model.output.Response;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@RestController
@RequestMapping("/chat")
public class ChatController {

//    OpenAiChatModel model;
    String apiKey;
    String baseUrl;

    public ChatController() {
        this.apiKey = System.getenv("API_KEY");
        this.baseUrl = System.getenv("BASE_URL");
//        this.model = OpenAiChatModel.builder()
//                //.baseUrl(baseUrl)
//                .apiKey("demo")
//                .modelName("gpt-4o-mini")
//                .build();
    }

    @GetMapping(value="/ai")
    String generation() {
        OpenAiChatModel model = OpenAiChatModel.builder()
                //.baseUrl(baseUrl)
                .apiKey("demo")
                .modelName("gpt-4o-mini")
                .build();
        String answer = model.generate("Say 'Hello World'");
        return answer;
    }

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE + ";charset=UTF-8")
    SseEmitter stream() {
        StreamingChatLanguageModel model = OpenAiStreamingChatModel.builder()
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .modelName("gpt-4o-mini")
                .build();

        String userMessage = "Tell me a joke";

        // 创建流式响应对象
        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
        model.generate(userMessage, new StreamingResponseHandler<AiMessage>() {

            @Override
            public void onNext(String token) {
                System.out.println("onNext: " + token);
                try {
                    emitter.send(token);
                } catch (IOException e) {
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onComplete(Response<AiMessage> response) {
                System.out.println("onComplete: " + response);
                emitter.complete();
            }

            @Override
            public void onError(Throwable error) {
                error.printStackTrace();
                emitter.complete();
            }
        });
        return emitter;
    }

    public static void main(String[] args) {
        System.setProperty("javax.net.ssl.trustStore", "/Users/chengyi/Documents/java_truestore/truststore.jks");	// 刚才填写的路径
        System.setProperty("javax.net.ssl.trustStorePassword", "chengyi");
        OpenAiStreamingLanguageModel model = OpenAiStreamingLanguageModel.builder()
                .apiKey("sk-67WVeNLNz4Icz6py6vR2SD8Bug3mVIAMSVQQt0eqiX879BuT")
                .baseUrl("https://www.dmxapi.cn/v1")
                .modelName("gpt-4o-mini")
                .build();

        String userMessage = "Tell me a joke";

        model.generate(userMessage, new StreamingResponseHandler<String>() {

            @Override
            public void onNext(String token) {
                System.out.println("onNext: " + token);
            }

            @Override
            public void onComplete(Response<String> response) {
                System.out.println("onComplete: " + response);
            }

            @Override
            public void onError(Throwable error) {
                error.printStackTrace();
            }
        });
    }

}
